package com.atguigu.ai.config;

import com.atguigu.ai.enums.SystemConstant;
import io.micrometer.observation.ObservationRegistry;
import io.netty.util.internal.UnstableApi;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.QdrantGrpcClient;
import io.qdrant.client.grpc.Collections;
import jakarta.annotation.PostConstruct;
import org.springframework.ai.autoconfigure.ollama.OllamaConnectionProperties;
import org.springframework.ai.autoconfigure.vectorstore.qdrant.QdrantVectorStoreProperties;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.openai.OpenAiEmbeddingOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.vectorstore.qdrant.QdrantVectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.concurrent.ExecutionException;

@Configuration
public class QdrantVectorConfig {

    @Autowired
    private QdrantVectorStoreProperties properties;

    @Bean
    public QdrantVectorStore openaiQdrantVectorStore(
            QdrantClient qdrantClient, OpenAiEmbeddingModel openAiEmbeddingModel
    ) throws ExecutionException, InterruptedException {
        if (!qdrantClient.collectionExistsAsync(SystemConstant.OPENAI_GPT3_QDRANT).get()) {
            qdrantClient.createCollectionAsync(SystemConstant.OPENAI_GPT3_QDRANT,
                    Collections.VectorParams.newBuilder()
                            .setDistance(Collections.Distance.Cosine).setSize(1536).build()).get();
        }
        return new QdrantVectorStore(
                qdrantClient,
                SystemConstant.OPENAI_GPT3_QDRANT,
                openAiEmbeddingModel,
                properties.isInitializeSchema()
        );
    }

    @Bean
    public QdrantVectorStore ollamaQdrantVectorStore(
            QdrantClient qdrantClient, OllamaEmbeddingModel ollamaEmbeddingModel
    ) throws ExecutionException, InterruptedException {
        if (!qdrantClient.collectionExistsAsync(SystemConstant.OLLAMA_QWEN2_QDRANT).get()) {
            qdrantClient.createCollectionAsync(SystemConstant.OLLAMA_QWEN2_QDRANT,
                    Collections.VectorParams.newBuilder()
                            .setDistance(Collections.Distance.Cosine).setSize(3584).build()).get();
        }
        return new QdrantVectorStore(
                qdrantClient,
                SystemConstant.OLLAMA_QWEN2_QDRANT,
                ollamaEmbeddingModel,
                properties.isInitializeSchema()
        );
    }

}
